{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp predefined_problems.cws_data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pre-defined Problems\n",
    "\n",
    "Preprocessing functions of pre-defined problems. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "\n",
    "import sys\n",
    "import os\n",
    "import glob\n",
    "from tqdm import tqdm\n",
    "import re\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "from m3tl.utils import filter_empty\n",
    "from m3tl.preproc_decorator import preprocessing_fn\n",
    "\n",
    "\n",
    "def process_line_msr_pku(l):\n",
    "    decoded_line = l.strip().split('  ')\n",
    "    return [w.strip('\\r\\n') for w in decoded_line]\n",
    "\n",
    "\n",
    "def process_line_as_training(l):\n",
    "    decoded_line = l.strip().split('\\u3000')\n",
    "    return [w.strip('\\r\\n') for w in decoded_line]\n",
    "\n",
    "\n",
    "def process_line_cityu(l):\n",
    "    decoded_line = l.strip().split(' ')\n",
    "    return [w.strip('\\r\\n') for w in decoded_line]\n",
    "\n",
    "\n",
    "def get_process_fn(filename):\n",
    "\n",
    "    if 'msr' in filename or 'pk' in filename:\n",
    "        return process_line_msr_pku\n",
    "\n",
    "    elif 'as' in filename:\n",
    "        return process_line_as_training\n",
    "\n",
    "    elif 'cityu' in filename:\n",
    "        return process_line_cityu\n",
    "\n",
    "\n",
    "def _process_text_files(path_list):\n",
    "\n",
    "    # Create possible tags for fast lookup\n",
    "    possible_tags = []\n",
    "    for i in range(1, 300):\n",
    "        if i == 1:\n",
    "            possible_tags.append('s')\n",
    "        else:\n",
    "            possible_tags.append('b' + 'm' * (i - 2) + 'e')\n",
    "\n",
    "    inputs = []\n",
    "    target = []\n",
    "\n",
    "    for s in range(len(path_list)):\n",
    "        filename = path_list[s]\n",
    "\n",
    "        # Init left and right queue\n",
    "\n",
    "        with open(filename, 'r', encoding='utf8') as f:\n",
    "\n",
    "            input_list = f.readlines()\n",
    "\n",
    "            process_fn = get_process_fn(os.path.split(filename)[-1])\n",
    "\n",
    "            for l in tqdm(input_list):\n",
    "\n",
    "                decoded_line = process_fn(l)\n",
    "\n",
    "                last_isalpha_num = False\n",
    "                pos_tag = []\n",
    "                final_line = []\n",
    "\n",
    "                for w in decoded_line:\n",
    "                    is_alphanum = bool(re.match('^[a-zA-Z0-9]+$', w))\n",
    "\n",
    "                    if not is_alphanum:\n",
    "                        # if this round is chinese and last round is alphanum,\n",
    "                        # append the last alphanum's target: s and continue\n",
    "                        if last_isalpha_num:\n",
    "                            pos_tag.append('s')\n",
    "\n",
    "                        if w and len(w) <= 299:\n",
    "                            final_line.extend(list(w))\n",
    "                            pos_tag.extend(list(possible_tags[len(w) - 1]))\n",
    "                        last_isalpha_num = False\n",
    "                    else:\n",
    "                        if last_isalpha_num:\n",
    "                            final_line[-1] += w\n",
    "                        else:\n",
    "                            final_line.append(w)\n",
    "                            last_isalpha_num = True\n",
    "\n",
    "                if last_isalpha_num:\n",
    "                    pos_tag.append('s')\n",
    "\n",
    "                # decode_str = ''.join(final_line)\n",
    "\n",
    "                # pos_tag_str = ''.join(pos_tag)\n",
    "\n",
    "                if len(pos_tag) != len(final_line):\n",
    "                    print(filename)\n",
    "                    print('Skip one row. ' + pos_tag + ';' + final_line)\n",
    "                    continue\n",
    "\n",
    "                inputs.append(final_line)\n",
    "                target.append(pos_tag)\n",
    "\n",
    "    return inputs, target\n",
    "\n",
    "\n",
    "def get_cws_fn(file_path):\n",
    "    @preprocessing_fn\n",
    "    def cws(params, mode):\n",
    "        file_list = glob.glob(file_path)\n",
    "\n",
    "        input_list = []\n",
    "        target_list = []\n",
    "\n",
    "        # Create possible tags for fast lookup\n",
    "        possible_tags = []\n",
    "        for i in range(1, 300):\n",
    "            if i == 1:\n",
    "                possible_tags.append('s')\n",
    "            else:\n",
    "                possible_tags.append('b' + 'm' * (i - 2) + 'e')\n",
    "\n",
    "        for file_path in file_list:\n",
    "            with open(file_path, 'r', encoding='utf8') as f:\n",
    "                raw_doc_list = f.readlines()\n",
    "            text_row_ind = [i+1 for i,\n",
    "                            text in enumerate(raw_doc_list) if '<S ID=' in text]\n",
    "\n",
    "            sentence_list = [text for i,\n",
    "                             text in enumerate(raw_doc_list) if i in text_row_ind]\n",
    "\n",
    "            for sentence in sentence_list:\n",
    "                input_list.append([])\n",
    "                target_list.append([])\n",
    "                for word in sentence.split():\n",
    "                    if word and len(word) <= 299:\n",
    "                        tag = possible_tags[len(word) - 1]\n",
    "                        input_list[-1] += list(word)\n",
    "                        target_list[-1] += list(tag)\n",
    "                    else:\n",
    "                        continue\n",
    "\n",
    "        if mode == 'train':\n",
    "            input_list, _, target_list, _ = train_test_split(\n",
    "                input_list, target_list, test_size=0.2, random_state=3721)\n",
    "        else:\n",
    "            _, input_list, _, target_list = train_test_split(\n",
    "                input_list, target_list, test_size=0.2, random_state=3721)\n",
    "\n",
    "        if mode == 'train':\n",
    "            file_list = [  # 'as_testing_gold.utf8',\n",
    "                'cityu_training.utf8', 'msr_training.utf8', 'pku_training.utf8']\n",
    "            file_list = [os.path.join('data/cws/training', f)\n",
    "                         for f in file_list]\n",
    "        else:\n",
    "            file_list = [  # 'as_testing_gold.utf8',\n",
    "                'cityu_test_gold.utf8', 'msr_test_gold.utf8', 'pku_test_gold.utf8']\n",
    "            # file_list = ['msr_test_gold.utf8']\n",
    "            file_list = [os.path.join('data/cws/gold', f) for f in file_list]\n",
    "\n",
    "        icwb_inputs, icwb_target = _process_text_files(file_list)\n",
    "\n",
    "        input_list += icwb_inputs\n",
    "        target_list += icwb_target\n",
    "\n",
    "        input_list, target_list = filter_empty(input_list, target_list)\n",
    "\n",
    "        return input_list, target_list\n",
    "    return cws\n",
    "\n",
    "\n",
    "def get_as_cws_fn(file_path):\n",
    "    @preprocessing_fn\n",
    "    def as_cws(params, mode):\n",
    "\n",
    "        if mode == 'train':\n",
    "            file_list = glob.glob('data/cws/training/as_*.utf8')\n",
    "        else:\n",
    "            file_list = ['as_testing_gold.utf8']\n",
    "            # file_list = ['msr_test_gold.utf8']\n",
    "            file_list = [os.path.join('data/cws/gold', f) for f in file_list]\n",
    "\n",
    "        input_list, target_list = _process_text_files(file_list)\n",
    "\n",
    "        return input_list, target_list\n",
    "    return as_cws\n",
    "\n",
    "\n",
    "def get_msr_cws_fn(file_path):\n",
    "    @preprocessing_fn\n",
    "    def msr_cws(params, mode):\n",
    "\n",
    "        if mode == 'train':\n",
    "            file_list = glob.glob('data/cws/training/msr_*.utf8')\n",
    "        else:\n",
    "            file_list = ['msr_test_gold.utf8']\n",
    "            # file_list = ['msr_test_gold.utf8']\n",
    "            file_list = [os.path.join('data/cws/gold', f) for f in file_list]\n",
    "\n",
    "        input_list, target_list = _process_text_files(file_list)\n",
    "\n",
    "        return input_list, target_list\n",
    "    return msr_cws\n",
    "\n",
    "\n",
    "def get_pku_cws_fn(file_path):\n",
    "    @preprocessing_fn\n",
    "    def pku_cws(params, mode):\n",
    "\n",
    "        if mode == 'train':\n",
    "            file_list = glob.glob('data/cws/training/pku_*.utf8')\n",
    "        else:\n",
    "            file_list = ['pku_test_gold.utf8']\n",
    "            # file_list = ['msr_test_gold.utf8']\n",
    "            file_list = [os.path.join('data/cws/gold', f) for f in file_list]\n",
    "\n",
    "        input_list, target_list = _process_text_files(file_list)\n",
    "\n",
    "        return input_list, target_list\n",
    "    return pku_cws\n",
    "\n",
    "\n",
    "def get_city_cws_fn(file_path):\n",
    "    @preprocessing_fn\n",
    "    def city_cws(params, mode):\n",
    "\n",
    "        if mode == 'train':\n",
    "            file_list = glob.glob('data/cws/training/cityu_*.utf8')\n",
    "        else:\n",
    "            file_list = ['cityu_test_gold.utf8']\n",
    "            # file_list = ['msr_test_gold.utf8']\n",
    "            file_list = [os.path.join('data/cws/gold', f) for f in file_list]\n",
    "\n",
    "        input_list, target_list = _process_text_files(file_list)\n",
    "\n",
    "        return input_list, target_list\n",
    "    return city_cws"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
